# import keras
import numpy as np
# import pandas as pd
import csv
import librosa
from keras.models import load_model

test_info_path = '/home/orient/文档/voiceprint_test_a/'


def get_mfcc(wav_file):
    wave, sr = librosa.load(test_info_path+'data/'+wav_file, sr=None)
    mfcc = librosa.feature.mfcc(wave, sr, n_mfcc=13)
    delta_mfcc = librosa.feature.delta(mfcc)
    delta2_mfcc = librosa.feature.delta(mfcc, order=2)
    feature = np.vstack([mfcc, delta_mfcc, delta2_mfcc])

    # reshape feature
    feature = mfcc_reshape(feature)

    return feature

def mfcc_reshape(*feature):
    mfcc = feature[0]
    zeros = np.zeros((mfcc.shape[0], 1))

    # slice
    while (mfcc.shape[1] < 100):
        mfcc = np.column_stack((mfcc, zeros))
    # intercept
    if (mfcc.shape[1] >= 100):
        mfcc = mfcc[:, :100]

    # Normalization
    mfcc = mfcc.reshape(mfcc.shape[0], -1) / np.max(mfcc)

    return mfcc


reg_list = []
with open(test_info_path+'reg.csv', 'r', encoding='UTF-8') as f:
    reader = csv.reader(f)
    headers = next(reader)
    for row in reader:
        reg_list.append(row[1])


try_list = []
with open(test_info_path+'try.csv', 'r', encoding='UTF-8') as f:
    reader = csv.reader(f)
    headers = next(reader)
    for row in reader:
        try_list.append(row[0])

model = load_model('/home/orient/PycharmProjects/OtSpeacherRecognization/model/CNN.h5')

i = 0

with open('/home/orient/PycharmProjects/OtSpeacherRecognization/res/res.csv', 'w') as f:
    writer = csv.writer(f)
    header = ["FILE_ID", "IS_FAMILY_MEMBER"]
    writer.writerow(header)
    for test in try_list:
        label = 'NO'
        i = i + 1
        j = 0
        for reg in reg_list:
            j = j + 1
            # print(test, reg)
            mfcc_test = get_mfcc(test+'.wav')
            mfcc_reg = get_mfcc(reg+'.wav')
            input = np.row_stack((mfcc_test, mfcc_reg))
            input = input.reshape(1, 78, 100)
            res = model.predict(input)
            print("i, j and res is: ", i, j, res[0][0])
            if(res[0][0] >= 0.98):
                label = 'YES'
                break
        content = [test, label]
        writer.writerow(content)
